# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/27 19:26 
@Author  : ZhangShenao 
@File    : train.py 
@Desc    : 训练Word2Vec模型
"""
from typing import List

from gensim.models import Word2Vec


def train_word2vec(sentences: List[List[str]]) -> None:
    """
    训练Word2Vec模型
    :param sentences: 与处理好的语料
    """

    # 使用gensim库,训练一个Word2Vec模型
    model = Word2Vec(sentences=sentences,  # 指定训练语料
                     vector_size=100,  # 设置词向量维度
                     window=5,  # 设置上下文窗口大小
                     min_count=1,  # 设置词频的最小阈值
                     workers=4,  # 设置进行训练的线程数
                     sg=1  # 模型架构 1:使用Skip-Gram架构 0:使用CBOW架构 这里因为是根据目标词预测上下文,因此采用Skip-Gram架构
                     )
    # 将模型保存到本地
    # word2vec.model: 主模型文件,包含了模型的参数、词汇表等信息。这个文件是加载完整模型所必需的
    # word2vec.model.wv.vectors.npy: 这个文件存储了模型中所有词汇的词向量
    # word2vec.model.syn1neg.npy: 这个文件存储的是训练过程中使用的负采样权重
    model.save("word2vec.model")

    print("Word2Vec 模型训练完成")
